/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
* This file is a part of the CANN Open Software.
* Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
* ===================================================================================================================*/
#ifndef GRAPH_ASCEND_REG_OPS_H
#define GRAPH_ASCEND_REG_OPS_H

#include <type_traits>
#include "graph/operator_reg.h"
#include "ascend_ir/ascend_ir.h"

namespace ge {
template <uint32_t OUTPUT_INDEX>
class AscOutputAttrDataType {
 public:
  AscOutputAttrDataType(ge::Operator *op) : op_(op) {}

  void operator=(const ge::DataType &value) {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(op_);
    const auto desc = ge::OpDescUtils::GetOpDescFromOperator(*op_);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc->MutableOutputDesc(OUTPUT_INDEX));
    desc->MutableOutputDesc(OUTPUT_INDEX)->SetDataType(value);
  }

  operator ge::DataType() const {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(op_);
    const auto desc = ge::OpDescUtils::GetOpDescFromOperator(*op_);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc->MutableOutputDesc(OUTPUT_INDEX));
    return desc->MutableOutputDesc(OUTPUT_INDEX)->GetDataType();
  };

 private:
  ge::Operator *op_;
};

template <uint32_t OUTPUT_INDEX>
class AscOutputAttrFormat {
 public:
  AscOutputAttrFormat(ge::Operator *op) : op_(op) {}

  void operator=(const ge::Format &value) {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(op_);
    const auto desc = ge::OpDescUtils::GetOpDescFromOperator(*op_);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc->MutableOutputDesc(OUTPUT_INDEX));
    desc->MutableOutputDesc(OUTPUT_INDEX)->SetFormat(value);
  }

  operator ge::Format() const {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(op_);
    const auto desc = ge::OpDescUtils::GetOpDescFromOperator(*op_);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(desc->MutableOutputDesc(OUTPUT_INDEX));
    return desc->MutableOutputDesc(OUTPUT_INDEX)->GetFormat();
  };

 private:
  ge::Operator *op_;
};

template <uint32_t OUTPUT_INDEX>
struct AscOpOutput {
 public:
  template <uint32_t INPUT_INDEX>
  friend class AscOpInput;

  AscOpOutput(ge::Operator *op) : op_(op), dtype(op), format(op),
                                  axis(AscTensorAttr::GetTensorAttr(op, OUTPUT_INDEX).axis),
                                  repeats(AscTensorAttr::GetTensorAttr(op, OUTPUT_INDEX).repeats),
                                  strides(AscTensorAttr::GetTensorAttr(op, OUTPUT_INDEX).strides) {}

  void SetContiguousView(const std::vector<Axis> &axes) {
    std::vector<AxisId> axes_ids;
    std::vector<ge::expression::ExpressionPtr> tmp_repeats;
    std::vector<ge::expression::ExpressionPtr> tmp_strides;
    axes_ids.reserve(axes.size());
    tmp_repeats.reserve(axes.size());
    tmp_strides.reserve(axes.size());

    std::for_each(axes.rbegin(), axes.rend(),
                  [&axes_ids, &tmp_repeats, &tmp_strides](const Axis &tmp_axis) {
                      if (tmp_strides.empty()) {
                          tmp_strides.emplace_back(ge::expression::Expression::CreateExpression(1));
                      } else {
                          tmp_strides.emplace_back(*tmp_repeats.rbegin() * *tmp_strides.rbegin());
                      }
                      tmp_repeats.emplace_back(tmp_axis.size);
                      axes_ids.emplace_back(tmp_axis.id);
                  });
    std::reverse(axes_ids.begin(), axes_ids.end());
    std::reverse(tmp_repeats.begin(), tmp_repeats.end());
    std::reverse(tmp_strides.begin(), tmp_strides.end());

    axis = axes_ids;
    repeats = tmp_repeats;
    strides = tmp_strides;
  }

  const ge::Operator &GetOwnerOp() {
    return *op_;
  }

 private:
  ge::Operator *op_;
 public:
  AscOutputAttrDataType<OUTPUT_INDEX> dtype;
  AscOutputAttrFormat<OUTPUT_INDEX> format;
  std::vector<int64_t> &axis;
  std::vector<ge::expression::ExpressionPtr> &repeats;
  std::vector<ge::expression::ExpressionPtr> &strides;
};

template<uint32_t INPUT_INDEX>
class AscOpInput {
 public:
  AscOpInput(ge::Operator *op) : op_(op) {}

  template<uint32_t OUTPUT_INDEX>
  AscOpInput<INPUT_INDEX> &operator=(const AscOpOutput<OUTPUT_INDEX> &output) {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(op_);
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(output.op_);
    op_->SetInput(INPUT_INDEX, *output.op_, OUTPUT_INDEX);
    AddEdgeForNode(*output.op_, OUTPUT_INDEX, *op_, INPUT_INDEX);
    return *this;
  }

 private:
  ge::Operator *op_;
};

template <const char* Name, typename T>
struct AscOpAttr {
 private:
  ge::Operator *op_;

 public:
  AscOpAttr(ge::Operator *op) : op_(op) {}
  void operator=(const T &value) {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(op_);
    this->op_->SetAttr(Name, value);
  }
};
}

#define OP_END_FACTORY_REG_WITHOUT_REGISTER(x) __OP_END_IMPL_WITHOUT_REGISTER__(x)

#endif  // GRAPH_ASCEND_REG_OPS_H
